package com.bw.test5;

import com.alibaba.alink.operator.batch.BatchOperator;
import com.alibaba.alink.operator.batch.source.MemSourceBatchOp;
import com.alibaba.alink.operator.common.evaluation.TuningRegressionMetric;
import com.alibaba.alink.params.shared.linear.LinearTrainParams;
import com.alibaba.alink.pipeline.regression.RidgeRegression;
import com.alibaba.alink.pipeline.tuning.GridSearchCV;
import com.alibaba.alink.pipeline.tuning.GridSearchCVModel;
import com.alibaba.alink.pipeline.tuning.ParamGrid;
import com.alibaba.alink.pipeline.tuning.RegressionTuningEvaluator;
import org.apache.flink.types.Row;
import org.junit.Test;

import java.util.Arrays;
import java.util.List;

public class GridSearchCVTest2 {
    @Test
    public void testRidgeRegression() throws Exception {
        BatchOperator.setParallelism(1);
        //1、构建数据集
        List<Row> df = Arrays.asList(
                Row.of(2, 1, 1),
                Row.of(3, 2, 1),
                Row.of(4, 3, 2),
                Row.of(2, 4, 1),
                Row.of(2, 2, 1),
                Row.of(4, 3, 2),
                Row.of(1, 2, 1)
        );
        BatchOperator<?> batchData = new MemSourceBatchOp(df, "f0 int, f1 int, label int");
        //2、创建岭回归pipline组件
        String[] colnames = new String[] {"f0", "f1"};
        RidgeRegression ridge = new RidgeRegression()
                .setFeatureCols(colnames)
                .setLambda(0.1)
                .setLabelCol("label")
                .setPredictionCol("pred");
        //3、调参
        ParamGrid paramGrid = new ParamGrid()
                .addGrid(ridge, RidgeRegression.LAMBDA,new Double[] {0.1,0.2})
                .addGrid(ridge, RidgeRegression.OPTIM_METHOD, new LinearTrainParams.OptimMethod[] {LinearTrainParams.OptimMethod.GD, LinearTrainParams.OptimMethod.SGD});

        // 回归评估
        RegressionTuningEvaluator tuningEvaluator = new RegressionTuningEvaluator()
                .setLabelCol("label")
                .setPredictionCol("pred")
                .setTuningRegressionMetric(TuningRegressionMetric.MSE);

        // 网格搜索
        GridSearchCV cv = new GridSearchCV()
                .setEstimator(ridge)
                .setParamGrid(paramGrid)
                .setTuningEvaluator(tuningEvaluator)
                .setNumFolds(2)
                .enableLazyPrintTrainInfo("TrainInfo");

        GridSearchCVModel model = cv.fit(batchData);
        //挑选最佳模型进行预测
        BatchOperator<?> result = model.getBestPipelineModel().transform(batchData);

        result.print();
    }
}
